# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import abc
import logging
from typing import List, Optional, Union

import tornado.gen
import requests
from recordclass import recordclass
import coin.proto.coin_order_enums_pb2 as coin_order

from coin.exchange.base.order_id_util import OrderIdGenerator
import coin.flow.subscriber as flow_subscriber
from coin.base.config import Config
from coin.base.timestamp import get_timestamp
from coin.flow.subscriber import TopicData, QueueData
from coin.exchange.base.kr_rest.currency import NativeCurrency
from coin.exchange.base.order_gateway import BaseOrderGatewayInfoMutator
from coin.proto.coin_query_pb2 import AccountBalance, ProductOrderElement
from coin.exchange.base.rest.async_private_client_base import AsyncPrivateClientBase
from coin.exchange.base.rest.private_queryer_base import PrivateQueryerBase
from coin.exchange.base.kr_rest.product import NativeProduct
from coin.exchange.base.order_gateway import (
    Order,
    OrderGatewayBase,
    OrderGatewayStatus,
    OrderSide,
    OrderType,
)

OrderInternal = recordclass('OrderInternal', [
    'exchange_order_id',
    'filled_qty',
    'timestamp',
])


def create_order_internal(exchange_order_id):
  return OrderInternal(exchange_order_id=exchange_order_id, filled_qty=None, timestamp=None)


def from_proto_order_side(order_side: int):
  order_side_map = {
      coin_order.BUY_ORDER: OrderSide.BUY,
      coin_order.SELL_ORDER: OrderSide.SELL,
  }
  if order_side in order_side_map:
    return order_side_map[order_side]
  else:
    raise ValueError("invalid order_side %s" % order_side)


def from_proto_order_type(order_type: int):
  order_type_map = {
      coin_order.DAY_LIMIT_ORDER: OrderType.LIMIT,
      coin_order.IOC_LIMIT_ORDER: OrderType.IOC,
  }
  if order_type in order_type_map:
    return order_type_map[order_type]
  else:
    raise ValueError("invalid order_type %s" % order_type)


def from_proto_order(proto_order: ProductOrderElement):
  internal = create_order_internal(proto_order.order_id)
  return Order(order_id=None,
               product=None,
               price=proto_order.price,
               qty=proto_order.qty,
               order_side=from_proto_order_side(proto_order.side),
               order_type=None,
               tag=None,
               accepted=False,
               cancel_sent=False,
               internal=internal)


class SpotGatewayInfoBase(BaseOrderGatewayInfoMutator, metaclass=abc.ABCMeta):
  BalanceType = float

  def __init__(self, currencies: List[NativeCurrency], products: List[NativeProduct]):
    super().__init__(currencies, products)
    assert products
    self._exchange = products[0].exchange

  @abc.abstractmethod
  def update_balance(self, balance_message):
    raise NotImplementedError()

  def balance_as_proto(self):
    return AccountBalance(exchange=self._exchange,
                          market_type="Spot",
                          each_balance=self.get_notNone_balances_internal())


class SpotOrderGatewayBase(OrderGatewayBase, metaclass=abc.ABCMeta):
  _market_type = "Spot"
  _exchange = None

  def __init__(self,
               config: Config,
               currencies,
               products: Union[List[NativeProduct], NativeProduct],
               logger=None):
    super().__init__()

    self._config = config

    # Orders
    self._working_orders = {}
    self._async_private_client = None
    self._order_id_generator = OrderIdGenerator()

    self._balance_ready = False
    self._orders_ready = False

    self._ioloop = None
    self._private_queryer = None

    # Flow
    self._flow_sub = None
    self._topic_strings = ['%s_rest_private_1' % self.exchange.lower()]

    self._logger = logger or logging.getLogger(__name__)

  def _check_ready(self):
    if self._balance_ready and self._orders_ready:
      self._og_ready.set_status(OrderGatewayStatus.READY)

  def start(self, queue, ioloop, flow_sub=None):
    self._async_private_client = AsyncPrivateClientBase(self._config.key_file)
    self._ioloop = ioloop

    # Prepare private queryer
    self._private_queryer = PrivateQueryerBase(queue,
                                               1,
                                               self._og_info.get_products(),
                                               self.exchange,
                                               self._async_private_client)
    self._private_queryer.start(ioloop)

    # Subscribe to flow for feed.
    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe(self._topic_strings[0], self.on_rest_private_msg)

  def stop(self):
    # Unsubscribe topic
    self._flow_sub.unsubscribe(self._topic_strings[0], self.on_rest_private_msg)
    self._private_queryer.stop()
    self._async_private_client.stop()

  def _handle_rest_private_balance(self, message):
    self._og_info.update_balance(message)
    self._balance_ready = True
    self._check_ready()

  @abc.abstractmethod
  def _handle_rest_private_list_orders(self, message):
    raise NotImplementedError()

  @abc.abstractmethod
  def _handle_rest_private_fills(self, message):
    raise NotImplementedError()

  def on_rest_private_msg(self, record, queue_data: QueueData, topic_data: TopicData):
    assert topic_data.topic_string
    assert queue_data.queue_name

    message = record.data

    query_type = message['query']
    if query_type == 'query_balance':
      self._handle_rest_private_balance(message['data'])
    elif query_type == 'query_list_orders':
      self._handle_rest_private_list_orders(message['data'])
    elif query_type == 'query_fills':
      self._handle_rest_private_fills(message['data'])
    else:
      raise ValueError('Unknown query: %s' % query_type)

  @tornado.gen.coroutine
  def submit_impl(
      self,
      product: NativeProduct,
      price: float,
      qty: float,
      order_side: OrderSide,
      order_type: OrderType = OrderType.LIMIT,  # LIMIT
      order_id: Optional[str] = None,
      tag=None,
      post_only=False):
    assert isinstance(product, NativeProduct)
    assert self.is_ready()
    assert price > 0
    assert qty > 0

    if order_side == OrderSide.BUY:
      side = 'bid'
    elif order_side == OrderSide.SELL:
      side = 'ask'
    else:
      raise ValueError('Unknown side: %s' % order_side)
    if order_type != OrderType.LIMIT:
      raise NotImplementedError()

    order_id = order_id or self._order_id_generator.gen_order_id()
    timestamp = get_timestamp()
    order = Order(order_id=order_id,
                  product=product,
                  price=price,
                  qty=qty,
                  order_side=order_side,
                  order_type=order_type,
                  tag=None,
                  accepted=False,
                  cancel_sent=False,
                  internal=OrderInternal(timestamp=timestamp,
                                         exchange_order_id=None,
                                         filled_qty=None))
    self._working_orders[order_id] = order

    try:
      update = yield self._async_private_client.async_query_submit(product, side, price, qty)
    except Exception as e:
      update = None
      self._logger.error('%s submit order error: %s', self.__class__.__name__, e)

    # After getting response.
    if update is not None and update.status == requests.codes.ok:
      exchange_order_id = self._extract_exchange_order_id(update.native_msg)
      order.internal.exchange_order_id = exchange_order_id
      order.accepted = True
      self._logger.info('%s submit order succeeded! %s', self.__class__.__name__, order_id)
    else:
      self._working_orders.pop(order_id, None)
      self._logger.info('%s submit order rejected! %s', self.__class__.__name__, order_id)

    elapsed_time = get_timestamp() - timestamp
    self._logger.debug('Order submit elapsed time:  %.4f sec', elapsed_time / 1e9)

  @abc.abstractmethod
  def _extract_exchange_order_id(self, resp_json):
    raise NotImplementedError()

  @tornado.gen.coroutine
  def cancel_impl(self, order_id, *, ignore_error):
    assert self.is_ready()

    if order_id in self._working_orders:
      if not self._working_orders[order_id].accepted:
        self._logger.error('Order has not been accepted: %s' % order_id)
        return

      if self._working_orders[order_id].cancel_sent:
        self._logger.error('Order has already cancel_sent: %s' % order_id)
        return

      order = self._working_orders[order_id]
      order.cancel_sent = True
      exchange_order_id = order.internal.exchange_order_id
      product = order.product
    else:
      if ignore_error:
        return
      else:
        raise ValueError('Invalid order id. %s' % order_id)

    timestamp = get_timestamp()
    try:
      update = yield self._async_private_client.async_query_cancel(product, exchange_order_id)
    except Exception as e:
      update = None
      self._logger.info('Cancel order error! %s' % e)

    if update is not None and update.status == requests.codes.ok:
      self._logger.info('Cancel order succeeded! %s' % order_id)
      if order:
        del self._working_orders[order_id]
    else:
      if order:
        order.cancel_sent = False
      self._logger.info('Cancel order rejected! %s' % order_id)

    elapsed_time = get_timestamp() - timestamp
    self._logger.info('Order cancel elapsed time:  %.4f sec', elapsed_time / 1e9)

  def cancel_multiple_impl(self, order_ids, *, ignore_error=False):
    for order_id in order_ids:
      self.cancel(order_id, ignore_error=ignore_error)

  def cancel_all_impl(self, *, ignore_error=False):
    for product in self._og_info.get_products():
      self._async_private_client.async_query_cancel_product(product)
    self._working_orders = {}

  def cancel_product_impl(self, product, *, ignore_error):
    order_ids = []
    for order_id, order in self._working_orders.items():
      if order.product == product:
        order_ids.append(order_id)
    self.cancel_multiple(order_ids, ignore_error=ignore_error)

  def get_working_order(self):
    return self._working_orders.values()

  def has_pending_order(self, product):
    for order in self._working_orders.values():
      if order.product == product and not order.accepted:
        return True
    return False

  def _find_order_by_exchange_order_id(self, exchange_order_id) -> Union[Order, None]:
    for order in self._working_orders.values():
      if order.internal.exchange_order_id == exchange_order_id:
        return order
    return None
